import cv2
import numpy as np
import transforms3d
import glob

class EyeinhandCalibration:
    def __init__(self, data_dir, board_size, square_size):
        self.images_dir = data_dir
        # 找棋盘格角点
        self.criteria_findCorner = (cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_FAST_CHECK + cv2.CALIB_CB_NORMALIZE_IMAGE)
        # 设置寻找亚像素角点的参数，采用的停止准则是最大循环次数30和最大误差容限0.001
        self.criteria_cornerSubPix = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001) # 阈值
        #棋盘格模板规格
        self.board_size = board_size
        # 世界坐标系中的棋盘格点,例如(0,0,0), (1,0,0), (2,0,0) ....,(8,5,0)，去掉Z坐标，记为二维矩阵
        objp = np.zeros((board_size[0]*board_size[1], 3), np.float32)
        objp[:,:2] = np.mgrid[0:board_size[1], 0:board_size[0]].T.reshape(-1,2) * square_size
        self.objp = objp # 10 mm

    def load_camera_param(self, camera_param_path: str):
        camera_param_cvfile = cv2.FileStorage(camera_param_path, cv2.FILE_STORAGE_READ)
        self.cameraMatrix = camera_param_cvfile.getNode("cameraMatrix").mat()
        self.distCoeffs = camera_param_cvfile.getNode("distCoeffs").mat()

    def load_position_info(self, pos_path: str):
        pose_vectors = []
        with open(pos_path, 'r') as fr:
            for line in fr.readlines():
                line = line.strip("\n")
                line_list = line.split(',')
                line_list = [float(i) for i in line_list]
                pose_vectors.append(line_list)
                # print(line_list)
        self.pose_vectors = np.array(pose_vectors)

    def pose_vectors_to_gripper2base_transforms(self, pose_vectors):
        # 提取旋转矩阵和平移向量
        R_gripper2bases = []
        t_gripper2bases = []
    
        # 迭代遍历每个位姿的旋转矩阵和平移向量
        for pose_vector in pose_vectors:
            # 提取旋转矩阵和平移向量
            R_gripper2base = self.rotation_vector_to_rotation_matrix(pose_vector[3], pose_vector[4], pose_vector[5])
            t_gripper2base = pose_vector[:3]
    
            # 提取旋转矩阵和平移向量
            R_gripper2bases.append(R_gripper2base)
            t_gripper2bases.append(t_gripper2base)
    
        return R_gripper2bases, t_gripper2bases
    
    def rotation_vector_to_rotation_matrix(self, rx, ry, rz):
        '''
        UR机械臂旋转向量表示, 
        将旋转向量转换为旋转矩阵
        :return: 旋转矩阵
        '''
        theta = np.sqrt(rx**2 + ry**2 + rz**2)
        x_rot_vector = rx / theta
        y_rot_vector = ry / theta
        z_rot_vector = rz / theta
    
        rotation_matrix = transforms3d.axangles.axangle2mat([x_rot_vector, y_rot_vector, z_rot_vector], theta)

        return rotation_matrix

    def eyeinhand_calibration(self, pos_path: str, camera_param_path: str):
        self.load_camera_param(camera_param_path)
        self.load_position_info(pos_path)
        images = glob.glob(self.images_dir + "/*.bmp")
        # 迭代处理图像
        obj_points = []  # 用于保存世界坐标系中的三维点
        img_points = []  # 用于保存图像平面上的二维点
        success_num = 0  # 用于保存检测成功的图像数量
        for i in range(len(images)):  # 遍历images:
            img = cv2.imread(self.images_dir + "/" + str(i) + '.bmp')   # 读取图像
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)   # RGB图像转换为灰度图像
            # 棋盘格检测
            ret, corners = cv2.findChessboardCorners(gray, (self.board_size[1], self.board_size[0]), self.criteria_findCorner)
            if ret:
                print("success_num: ", success_num)
                success_num += 1
                corners = cv2.cornerSubPix(gray, corners, (11, 11), (-1, -1), self.criteria_cornerSubPix)

                # 如果成功检测到棋盘格，添加图像平面上的二维点和世界坐标系中的三维点到列表
                obj_points.append(self.objp)
                img_points.append(corners)

                # 绘制并显示角点
                cv2.drawChessboardCorners(img, (self.board_size[1], self.board_size[0]), corners, ret)
                # cv2.imwrite("imgtest.bmp", img)
                cv2.circle(img, (int(corners[1, 0, 0]), int(corners[1, 0, 1])), 4, (0,0,255), 10)
                cv2.namedWindow('img', cv2.WINDOW_NORMAL)
                cv2.resizeWindow('img', 640, 480)
                cv2.imshow('img', img)
                cv2.waitKey(200)

        cv2.destroyAllWindows()

        # # 打印obj_point和img_point的形状
        print(np.array(obj_points).shape)
        print(np.array(img_points).shape)

        # 求解标定板位姿
        R_target2cams = []  # 用于保存旋转矩阵
        t_target2cams = []  # 用于保存平移向量
        # 迭代的到每张图片相对于相机的位姿
        for i in range(success_num):
            # rvec：标定板相对于相机坐标系的旋转向量
            # t_target2cam：标定板相对于相机坐标系的平移向量
            ret, rvec_target2cam, t_target2cam = cv2.solvePnP(obj_points[i], img_points[i], self.cameraMatrix, self.distCoeffs) 

            # 将旋转向量(rvec)转换为旋转矩阵
            # R_target2cam：标定板相对于相机坐标系的旋转矩阵
            R_target2cam, _ = cv2.Rodrigues(rvec_target2cam)   # 输出：R为旋转矩阵和旋转向量的关系  输入：rvec为旋转向量

            # 将标定板相对于相机坐标系的旋转矩阵和平移向量保存到列表
            R_target2cams.append(R_target2cam)
            t_target2cams.append(t_target2cam)

        # 求解手眼标定
        R_gripper2bases, t_gripper2bases = self.pose_vectors_to_gripper2base_transforms(self.pose_vectors)

        # R_camera2end：相机相对于机械臂末端的旋转矩阵
        # t_camera2end：相机相对于机械臂末端的平移向量
        R_cam2gripper, t_cam2gripper = cv2.calibrateHandEye(R_gripper2bases, t_gripper2bases, R_target2cams, t_target2cams, method=cv2.CALIB_HAND_EYE_TSAI)
        fs = cv2.FileStorage('eyeinhand_params.yaml', cv2.FILE_STORAGE_WRITE)
        fs.write("R_cam2gripper", R_cam2gripper)
        fs.write("t_cam2gripper", t_cam2gripper)
        fs.release()
        
        # 将旋转矩阵和平移向量组合成齐次位姿矩阵
        T_camera2end = np.eye(4)
        T_camera2end[:3, :3] = R_cam2gripper
        T_camera2end[:3, 3] = t_cam2gripper.reshape(3)


        # 输出相机相对于机械臂末端的旋转矩阵和平移向量
        print("Camera to end rotation matrix:")
        print(R_cam2gripper)
        print("Camera to end translation vector:") 
        print(t_cam2gripper)

        # 输出相机相对于机械臂末端的位姿矩阵
        print("Camera to end pose matrix:")
        np.set_printoptions(suppress=True)  # suppress参数用于禁用科学计数法
        print(T_camera2end)

        


if __name__ == "__main__":
        images_dir = "D:/VSCodeProjects/calibration/data/eyeinhand20250427"
        calib = EyeinhandCalibration(images_dir, board_size=(8,11), square_size=10.0)

        pos_path = "D:/VSCodeProjects/calibration/data/eyeinhand20250427.txt"
        camera_param_path = "D:/VSCodeProjects/calibration/camera_params.yaml"
        calib.eyeinhand_calibration(pos_path, camera_param_path)